import numpy as np
import matplotlib.pyplot as plt

from globals import *


'''
气动力（电压值）滤波
'''
def filter_aero(input_data, freq_sample,freq_motion):
    
    input_data = np.array(input_data)
    
    filtered_list = np.zeros((len(input_data),len(input_data[0])))
    filtered_list[:,-1] = input_data[:,-1]

    # plt.figure(figsize=FIG_SIZE)
    for i in range(12):
        noisy_signal = input_data[:,i]

        # 进行傅里叶变换
        fft_result = np.fft.fft(noisy_signal)

        # 构建频率轴
        frequencies = np.fft.fftfreq(len(fft_result),1/freq_sample)

        # 滤波（这里假设要去除高频成分）
        if(max(freq_motion))>0:
            cutoff = 2*max(freq_motion)     # 10倍运动周期（默认4个运动周期）
        else:
            cutoff = 0.2
        # cutoff = 0.2

        fft_result_filtered = fft_result.copy()
        fft_result_filtered[np.abs(frequencies) > cutoff] = 0
        # fft_result_filtered[np.abs(frequencies) < 0.001] = 0


        # 反变换回时域
        filtered_signal = np.fft.ifft(fft_result_filtered)

        filtered_list[:,i] = np.real(filtered_signal.transpose())

        

        freq_range = 10  # 设置要显示的频率范围为 0~10Hz
        indices = np.where((frequencies >= 0) & (frequencies <= freq_range))[0]

    #     plt.subplot(4,3,i+1)
    #     plt.title("u"+str(i+1))
    #     plt.plot(frequencies[indices], np.abs(fft_result[indices]),label='Original')
    #     plt.plot(frequencies[indices], np.abs(fft_result_filtered[indices]),label='Filtered')
    #     plt.xlabel('Frequency')
    #     plt.ylabel('Amplitude')
    #     plt.legend()
    # plt.show()


    print("截止频率："+str(cutoff)+"Hz")
    
    return filtered_list


'''
绳拉力滤波
'''
def filter_tension(input_data, freq):
    pass

'''
解码器滤波
'''
def filter_encoder(input_data, freq):
    pass

'''
陀螺仪数据滤波
'''
def filter_angle(input_data):
    input_data=np.array(input_data)
    yaw0=input_data[0,2]
    # yaw = np.average(input_data[:,2])
    input_data[:,2]=input_data[:,2]-yaw0
    return input_data

'''
气动力数据滤波
'''
def filter_froce(input_data):
    input_data=np.array(input_data)

    # length = len(input_data[:,1]) // 2

    # fy = np.average(input_data[:,1])
    # fy = np.average(input_data[length//2,1])
    fy = (np.max(input_data[:,1])+np.min(input_data[:,1]))/2.0
    input_data[:,1]=input_data[:,1]-fy

    # mx = np.average(input_data[:,3])
    # mx = np.average(input_data[length//2,1])
    mx = (np.max(input_data[:,3])+np.min(input_data[:,3]))/2.0
    input_data[:,3]=input_data[:,3]-mx

    # mz = np.average(input_data[:,5])
    # mz = np.average(input_data[length//2,1])
    mz = (np.max(input_data[:,5])+np.min(input_data[:,5]))/2.0
    input_data[:,5]=input_data[:,5]-mz


    return input_data
